#include "viml/parser.h"

namespace viml {

Parser::Parser(const std::vector<Token>& tokens) : tokens(tokens), current(0) {}

Parser::Parser(const std::string& source) : current(0) {
    Lexer lexer(source);
    tokens = lexer.tokenize();
}

Token Parser::peek() const {
    if (current < tokens.size()) {
        return tokens[current];
    }
    return Token(TokenType::END_OF_FILE, "", 0, 0);
}

Token Parser::previous() const {
    if (current > 0 && current <= tokens.size()) {
        return tokens[current - 1];
    }
    return Token(TokenType::END_OF_FILE, "", 0, 0);
}

bool Parser::is_at_end() const {
    return current >= tokens.size() || peek().type == TokenType::END_OF_FILE;
}

bool Parser::check(TokenType type) const {
    if (is_at_end()) return false;
    return peek().type == type;
}

bool Parser::check(const std::vector<TokenType>& types) const {
    if (is_at_end()) return false;
    for (TokenType type : types) {
        if (peek().type == type) {
            return true;
        }
    }
    return false;
}

bool Parser::match(TokenType type) {
    if (check(type)) {
        advance();
        return true;
    }
    return false;
}

bool Parser::match(const std::vector<TokenType>& types) {
    for (TokenType type : types) {
        if (check(type)) {
            advance();
            return true;
        }
    }
    return false;
}

Token Parser::advance() {
    if (!is_at_end()) {
        current++;
    }
    return previous();
}

Token Parser::consume(TokenType type, const std::string& message) {
    if (check(type)) {
        return advance();
    }
    Token token = peek();
    throw ParseError(message, token.line, token.column);
}

void Parser::skip_newlines() {
    while (match(TokenType::NEWLINE)) {
        // 跳过换行符
    }
}

std::unique_ptr<Program> Parser::parse() {
    std::vector<std::unique_ptr<Statement>> statements;
    
    skip_newlines();
    while (!is_at_end()) {
#ifdef DEBUG_BUILD
        std::cout << "Parsing statement, current token: " << static_cast<int>(peek().type) << std::endl;
#endif
        try {
            auto stmt = statement();
            if (stmt) {
                statements.push_back(std::move(stmt));
            }
            skip_newlines();
        } catch (const ParseError& e) {
            // 错误恢复
            synchronize();
        }
    }
    
    return std::make_unique<Program>(std::move(statements));
}

std::unique_ptr<Expression> Parser::parse_expression() {
    return expression();
}

std::unique_ptr<Statement> Parser::statement() {
    skip_newlines();
    
    if (match(TokenType::LET)) {
        return let_statement();
    }
    if (match(TokenType::ECHO)) {
        return echo_statement();
    }
    if (match(TokenType::IF)) {
        return if_statement();
    }
    if (match(TokenType::WHILE)) {
        return while_statement();
    }
    if (match(TokenType::FOR)) {
        return for_statement();
    }
    if (match(TokenType::FUNCTION)) {
        return function_statement();
    }
    if (match(TokenType::RETURN)) {
        return return_statement();
    }
    if (match(TokenType::CALL)) {
#ifdef DEBUG_BUILD
        std::cout << "Parsing CALL statement" << std::endl;
#endif
        // 解析函数调用表达式
        auto expr = expression();
        // 确保这是一个CallExpression
        if (dynamic_cast<CallExpression*>(expr.get())) {
#ifdef DEBUG_BUILD
            std::cout << "Detected CallExpression in CALL statement" << std::endl;
#endif
            // 转移所有权
            auto call_ptr = std::unique_ptr<CallExpression>(static_cast<CallExpression*>(expr.release()));
            return std::make_unique<CallStatement>(std::move(call_ptr));
        } else {
            error("Expected function call after 'call' keyword");
        }
    }
    
    // 表达式语句（函数调用等）
    auto expr = expression();
#ifdef DEBUG_BUILD
    std::cout << "Expression type: " << typeid(*expr).name() << std::endl;
#endif
    // 检查是否是函数调用表达式
    if (dynamic_cast<CallExpression*>(expr.get())) {
#ifdef DEBUG_BUILD
        std::cout << "Detected CallExpression in statement" << std::endl;
#endif
        // 转移所有权
        auto call_ptr = std::unique_ptr<CallExpression>(static_cast<CallExpression*>(expr.release()));
        return std::make_unique<CallStatement>(std::move(call_ptr));
    } else {
#ifdef DEBUG_BUILD
        std::cout << "Creating EchoStatement for expression" << std::endl;
#endif
        // 其他表达式作为Echo处理
        // 使用表达式中第一个token的行号
        size_t stmt_line = 0;
        size_t stmt_column = 0;
        if (expr) {
            stmt_line = expr->line;
            stmt_column = expr->column;
        }
        return std::make_unique<EchoStatement>(std::move(expr), stmt_line, stmt_column);
    }
}

std::unique_ptr<Statement> Parser::let_statement() {
    Token name = consume(TokenType::IDENTIFIER, "Expected variable name");
    consume(TokenType::ASSIGN, "Expected '=' after variable name");
    auto value = expression();
    return std::make_unique<LetStatement>(name.value, std::move(value), name.line, name.column);
}

// 辅助函数：解析一个完整的表达式序列，正确处理运算符优先级
std::unique_ptr<Expression> Parser::parse_full_expression() {
    return expression();
}

std::unique_ptr<Statement> Parser::echo_statement() {
    std::vector<std::unique_ptr<Expression>> expressions;
    
    // 解析第一个表达式
    expressions.push_back(parse_full_expression());
    
    // 解析后续的表达式（用逗号或空格分隔）
    while (!check(TokenType::NEWLINE) && !is_at_end()) {
        // 检查是否是逗号分隔
        if (match(TokenType::COMMA)) {
            expressions.push_back(parse_full_expression());
        } 
        // 检查是否是空格分隔（即下一个token是表达式开始）
        else if (is_expression_start()) {
            expressions.push_back(parse_full_expression());
        } 
        // 其他情况退出循环
        else {
            break;
        }
    }
    
    // 如果只有一个表达式，直接使用
    if (expressions.size() == 1) {
        return std::make_unique<EchoStatement>(std::move(expressions[0]));
    }
    
    // 如果有多个表达式，创建一个列表字面量
    auto list_literal = std::make_unique<ListLiteral>(std::move(expressions));
    return std::make_unique<EchoStatement>(std::move(list_literal));
}

std::unique_ptr<Statement> Parser::if_statement() {
    auto condition = expression();
    skip_newlines();
    
    std::vector<std::unique_ptr<Statement>> then_body = block();
    std::vector<std::unique_ptr<Statement>> else_body;
    
    if (match(TokenType::ELSEIF)) {
        // elseif转换为嵌套的if语句递归处理
        auto nested_if = if_statement();
        else_body.push_back(std::move(nested_if));
    } else if (match(TokenType::ELSE)) {
        skip_newlines();
        else_body = block();
        consume(TokenType::ENDIF, "Expected 'endif'");
    } else {
        consume(TokenType::ENDIF, "Expected 'endif'");
    }
    
    return std::make_unique<IfStatement>(std::move(condition), std::move(then_body), std::move(else_body));
}

std::unique_ptr<Statement> Parser::while_statement() {
    auto condition = expression();
    skip_newlines();
    
    std::vector<std::unique_ptr<Statement>> body = block();
    
    consume(TokenType::ENDWHILE, "Expected 'endwhile'");
    return std::make_unique<WhileStatement>(std::move(condition), std::move(body));
}

std::unique_ptr<Statement> Parser::for_statement() {
    Token variable = consume(TokenType::IDENTIFIER, "Expected variable name");
    consume(TokenType::IN, "Expected 'in' after variable");
    auto iterable = expression();
    skip_newlines();
    
    std::vector<std::unique_ptr<Statement>> body = block();
    
    consume(TokenType::ENDFOR, "Expected 'endfor'");
    return std::make_unique<ForStatement>(variable.value, std::move(iterable), std::move(body));
}

std::unique_ptr<Statement> Parser::function_statement() {
    Token name = consume(TokenType::IDENTIFIER, "Expected function name");
    consume(TokenType::LPAREN, "Expected '(' after function name");
    
    std::vector<std::string> parameters;
    if (!check(TokenType::RPAREN)) {
        do {
            Token param = consume(TokenType::IDENTIFIER, "Expected parameter name");
            parameters.push_back(param.value);
        } while (match(TokenType::COMMA));
    }
    consume(TokenType::RPAREN, "Expected ')' after parameters");
    skip_newlines();
    
    std::vector<std::unique_ptr<Statement>> body = block();
    
    consume(TokenType::ENDFUNCTION, "Expected 'endfunction'");
    return std::make_unique<FunctionStatement>(name.value, std::move(parameters), std::move(body));
}

std::unique_ptr<Statement> Parser::return_statement() {
    std::unique_ptr<Expression> value = nullptr;
    if (!check(TokenType::NEWLINE) && !is_at_end()) {
        value = expression();
    }
    return std::make_unique<ReturnStatement>(std::move(value));
}

std::vector<std::unique_ptr<Statement>> Parser::block() {
    std::vector<std::unique_ptr<Statement>> statements;
    
    while (!is_at_end() && 
           !check(TokenType::ENDIF) && 
           !check(TokenType::ELSEIF) &&
           !check(TokenType::ELSE) && 
           !check(TokenType::ENDWHILE) && 
           !check(TokenType::ENDFOR) && 
           !check(TokenType::ENDFUNCTION)) {
        try {
            auto stmt = statement();
            if (stmt) {
                statements.push_back(std::move(stmt));
            }
            skip_newlines();
        } catch (const ParseError& e) {
            synchronize();
        }
    }
    
    return statements;
}

std::unique_ptr<Expression> Parser::expression() {
    return logical_or();
}

std::unique_ptr<Expression> Parser::logical_or() {
    auto expr = logical_and();
    
    while (match(TokenType::OR)) {
        auto right = logical_and();
        expr = std::make_unique<BinaryExpression>(std::move(expr), BinaryExpression::Operator::OR, std::move(right), previous().line, previous().column);
    }
    
    return expr;
}

std::unique_ptr<Expression> Parser::logical_and() {
    auto expr = equality();
    
    while (match(TokenType::AND)) {
        auto right = equality();
        expr = std::make_unique<BinaryExpression>(std::move(expr), BinaryExpression::Operator::AND, std::move(right), previous().line, previous().column);
    }
    
    return expr;
}

std::unique_ptr<Expression> Parser::equality() {
    auto expr = comparison();
    
    while (match({TokenType::EQUAL, TokenType::NOT_EQUAL})) {
        Token op = previous();
        auto right = comparison();
        BinaryExpression::Operator bin_op = (op.type == TokenType::EQUAL) ? 
            BinaryExpression::Operator::EQ : BinaryExpression::Operator::NE;
        expr = std::make_unique<BinaryExpression>(std::move(expr), bin_op, std::move(right), op.line, op.column);
    }
    
    return expr;
}

std::unique_ptr<Expression> Parser::comparison() {
    auto expr = term();
    
    while (match({TokenType::GREATER_THAN, TokenType::GREATER_EQUAL, 
                  TokenType::LESS_THAN, TokenType::LESS_EQUAL})) {
        Token op = previous();
        auto right = term();
        
        BinaryExpression::Operator bin_op;
        switch (op.type) {
            case TokenType::GREATER_THAN: bin_op = BinaryExpression::Operator::GT; break;
            case TokenType::GREATER_EQUAL: bin_op = BinaryExpression::Operator::GE; break;
            case TokenType::LESS_THAN: bin_op = BinaryExpression::Operator::LT; break;
            case TokenType::LESS_EQUAL: bin_op = BinaryExpression::Operator::LE; break;
            default: throw ParseError("Invalid comparison operator", op.line, op.column);
        }
        
        expr = std::make_unique<BinaryExpression>(std::move(expr), bin_op, std::move(right), op.line, op.column);
    }
    
    return expr;
}

std::unique_ptr<Expression> Parser::term() {
    auto expr = factor();
    
    while (match({TokenType::PLUS, TokenType::MINUS, TokenType::DOT})) {
        Token op = previous();
        auto right = factor();
        BinaryExpression::Operator bin_op;
        if (op.type == TokenType::PLUS) {
            bin_op = BinaryExpression::Operator::ADD;
        } else if (op.type == TokenType::MINUS) {
            bin_op = BinaryExpression::Operator::SUB;
        } else if (op.type == TokenType::DOT) {
            bin_op = BinaryExpression::Operator::ADD;  // DOT在VimL中表示字符串连接，等同于ADD
        } else {
            throw ParseError("Invalid term operator", op.line, op.column);
        }
        expr = std::make_unique<BinaryExpression>(std::move(expr), bin_op, std::move(right), op.line, op.column);
    }
    
    return expr;
}

std::unique_ptr<Expression> Parser::factor() {
    auto expr = unary();
    
    while (match({TokenType::MULTIPLY, TokenType::DIVIDE, TokenType::MODULO})) {
        Token op = previous();
        auto right = unary();
        
        BinaryExpression::Operator bin_op;
        switch (op.type) {
            case TokenType::MULTIPLY: bin_op = BinaryExpression::Operator::MUL; break;
            case TokenType::DIVIDE: bin_op = BinaryExpression::Operator::DIV; break;
            case TokenType::MODULO: bin_op = BinaryExpression::Operator::MOD; break;
            default: throw ParseError("Invalid factor operator", op.line, op.column);
        }
        
        expr = std::make_unique<BinaryExpression>(std::move(expr), bin_op, std::move(right), op.line, op.column);
    }
    
    return expr;
}

std::unique_ptr<Expression> Parser::unary() {
    if (match({TokenType::NOT, TokenType::MINUS})) {
        Token op = previous();
        auto right = unary();
        UnaryExpression::Operator un_op = (op.type == TokenType::NOT) ? 
            UnaryExpression::Operator::NOT : UnaryExpression::Operator::MINUS;
        return std::make_unique<UnaryExpression>(un_op, std::move(right));
    }
    
    return call();
}

std::unique_ptr<Expression> Parser::call() {
    auto expr = primary();
    
    while (true) {
        if (match(TokenType::LPAREN)) {
            std::vector<std::unique_ptr<Expression>> arguments;
            if (!check(TokenType::RPAREN)) {
                do {
                    arguments.push_back(expression());
                } while (match(TokenType::COMMA));
            }
            consume(TokenType::RPAREN, "Expected ')' after arguments");
            expr = std::make_unique<CallExpression>(std::move(expr), std::move(arguments));
        } else if (match(TokenType::LBRACKET)) {
            auto index = expression();
            consume(TokenType::RBRACKET, "Expected ']' after index");
            expr = std::make_unique<IndexExpression>(std::move(expr), std::move(index));
        } else {
            break;
        }
    }
    
    return expr;
}

std::unique_ptr<Expression> Parser::primary() {
    if (match(TokenType::NUMBER)) {
        double value = std::stod(previous().value);
        return std::make_unique<NumberLiteral>(value, previous().line, previous().column);
    }

    if (match(TokenType::STRING)) {
        return std::make_unique<StringLiteral>(previous().value, previous().line, previous().column);
    }

    if (match(TokenType::IDENTIFIER)) {
        return std::make_unique<Identifier>(previous().value, previous().line, previous().column);
    }

    if (match(TokenType::SPECIAL_CONSTANT)) {
        return std::make_unique<SpecialConstantNode>(previous().value, previous().line, previous().column);
    }
    
    if (match(TokenType::LPAREN)) {
        auto expr = expression();
        consume(TokenType::RPAREN, "Expected ')' after expression");
        return expr;
    }
    
    if (match(TokenType::LBRACKET)) {
        return list_literal();
    }
            
    if (match(TokenType::LBRACE)) {
        return dict_literal();
    }
    
    Token token = peek();
    throw ParseError("Unexpected token: " + token.value, token.line, token.column);
}

std::unique_ptr<Expression> Parser::list_literal() {
    std::vector<std::unique_ptr<Expression>> elements;
    
    if (!check(TokenType::RBRACKET)) {
        do {
            elements.push_back(expression());
        } while (match(TokenType::COMMA));
    }
    
    consume(TokenType::RBRACKET, "Expected ']' after list elements");
    return std::make_unique<ListLiteral>(std::move(elements));
}

std::unique_ptr<Expression> Parser::dict_literal() {
    std::vector<std::pair<std::unique_ptr<Expression>, std::unique_ptr<Expression>>> pairs;
    
    // 在检查空字典之前，跳过可能的换行符
    while (check(TokenType::NEWLINE)) {
        advance();
    }
    
    if (!check(TokenType::RBRACE)) {
        do {
            // 在解析每个键之前跳过换行符
            while (check(TokenType::NEWLINE)) {
                advance();
            }
            
            auto key = expression();
            consume(TokenType::COLON, "Expected ':' after dictionary key");
            auto value = expression();
            pairs.emplace_back(std::move(key), std::move(value));
            
            // 在检查逗号之前跳过换行符
            while (check(TokenType::NEWLINE)) {
                advance();
            }
        } while (match(TokenType::COMMA));
    }
    
    // 在检查结束括号之前跳过换行符
    while (check(TokenType::NEWLINE)) {
        advance();
    }
    
    consume(TokenType::RBRACE, "Expected '}' after dictionary elements");
    return std::make_unique<DictLiteral>(std::move(pairs));
}

void Parser::error(const std::string& message) {
    Token token = peek();
    throw ParseError(message, token.line, token.column);
}

void Parser::synchronize() {
    advance();
    
    while (!is_at_end()) {
        if (previous().type == TokenType::NEWLINE) return;
        
        switch (peek().type) {
            case TokenType::LET:
            case TokenType::IF:
            case TokenType::WHILE:
            case TokenType::FOR:
            case TokenType::FUNCTION:
            case TokenType::RETURN:
            case TokenType::ECHO:
                return;
            default:
                break;
        }
        
        advance();
    }
}

// 辅助函数：检查下一个token是否是表达式的开始
bool Parser::is_expression_start() {
    return check({TokenType::NUMBER, TokenType::STRING, TokenType::IDENTIFIER, 
                  TokenType::LPAREN, TokenType::LBRACKET, TokenType::LBRACE,
                  TokenType::NOT, TokenType::SPECIAL_CONSTANT});  // 添加特殊常量支持
}

} // namespace viml